#code up Fst analysis
#relies on some data generated from MKtests.R

grouse.df2 <- data.frame(cDNA=character(0),id=character(0),start.pos=numeric(0),pop.1.freq=numeric(0),pop.2.freq=numeric(0),pop.3.freq=numeric(0),pop.1.tsv.depth=numeric(0),pop.2.tsv.depth=numeric(0),pop.3.tsv.depth=numeric(0),Fst=numeric(0),stringsAsFactors=FALSE)
for(i in 1:length(grouseFst.df)){
	if(is.null(grouseFst.df[[i]])) next
	grouse.df2 <- rbind(grouse.df2,grouseFst.df[[i]])
	}
grouse.df2[,c("pop.1.count","pop.2.count","pop.3.count")] <- round(grouse.df2[,c("pop.1.freq","pop.2.freq","pop.3.freq")] * grouse.df2[,c("pop.1.tsv.depth","pop.2.tsv.depth","pop.3.tsv.depth")])
#looks ok

grouse.df2$i.vec <- as.numeric(as.factor(grouse.df2$cDNA))
grouse.df2$k.vec <- 0

for(i in 1:max(grouse.df2$i.vec)){
	grouse.df2$k.vec[grouse.df2$i.vec==i] <- 1:sum(grouse.df2$i.vec==i) 
	}

Fst.data1 <- list(counts = as.matrix(grouse.df2[,c("pop.1.count","pop.2.count","pop.3.count")]),
	depth = as.matrix(grouse.df2[,c("pop.1.tsv.depth","pop.2.tsv.depth","pop.3.tsv.depth")]),
	i.vec = grouse.df2$i.vec, k.vec = grouse.df2$k.vec, IK = nrow(grouse.df2),J=3,
	K=as.vector(by(grouse.df2$k.vec,grouse.df2$i.vec,max))
	)

Fst.data2 <- list(counts=Fst.data1$counts[apply(Fst.data1$depth,1,function(x){all(x>0)}),],
	depth = Fst.data1$depth[apply(Fst.data1$depth,1,function(x){all(x>0)}),],
	i.vec = grouse.df2$i.vec[apply(Fst.data1$depth,1,function(x){all(x>0)})],J=3)
Fst.data2$IK <- nrow(Fst.data2$counts)


#set initial values
Fst.init1 <- list(eta = matrix(-1.51,nrow=nrow(Fst.data1$count),ncol=3))
Fst.init2 <- list(
	p=Fst.data2$counts/Fst.data2$depth) #eta = matrix(-1.51,nrow=nrow(Fst.data2$count),ncol=3),
require(rjags)

tst <- jags.model('bug_model3.txt',data=Fst.data2,inits=Fst.init2)
tst.coda <- coda.samples(tst,c("alpha","beta"),n.iter=10,thin=1)


jm.fst2 <- jags.model('bug_model3.txt',data=Fst.data2,inits=Fst.init2,n.adapt=5000)
jm.fst2.coda <- coda.samples(jm.fst2,c("alpha","beta","F","a","tau.ijk"),n.iter=10000,thin=20)

#################
#jags not working
#################

#read data in for Beaumont and Balding
grouse.df3 <- grouse.df2[apply(grouse.df2[,grep('depth',names(grouse.df2))],1,function(x){all(x>0)}),]

writeBayesFst <- function(data,file,gamma=1){
	count.cols <- grep('count',names(data))
	depth.cols <- grep('depth',names(data))
	if(length(count.cols)!=length(depth.cols)) stop('must have equal number of counts and depths')
	cat(length(count.cols)," populations found\n")
	
	cat(as.character(gamma),"\n",file=file)
	cat(as.character(length(count.cols)),"\n",file=file,append=TRUE) #no. populations
	cat(as.character(nrow(data)),"\n",file=file,append=TRUE) #no. loci
	
	for(i in 1:nrow(data)){
		cat(paste("2","%",data[i,"id"],sep="\t"),"\n",file=file,append=TRUE)
		for(j in 1:length(count.cols)){
			count.j <- c(data[i,count.cols[j]], data[i,depth.cols[j]]-data[i,count.cols[j]])
			count.j[count.j<0] <- 0
			cat(paste(count.j,collapse="\t"),"\n",file=file,append=TRUE)
			}
		}
	}


writeBayesFst(grouse.df3, "../test_files/BB_Fst.inp")
writeBayesFst(grouse.df3, "../test_files/BB_Fst0.inp",gamma=0)

#problems with convergence??
writeBayesFst(grouse.df3[seq(2,nrow(grouse.df3),2),], "../test_files/BB_Fst0_1.inp",gamma=0)

grouse.df3$p.hat <- apply(grouse.df3[,grep('freq',names(grouse.df3))],1,mean)
grouse.df3$mean.het <- -1

calculateHeterozygosity <- function(X){
	weights <- X[,grep("depth",names(X))]
	total.depth <- sum(weights)
	p <- X[,grep("freq",names(X))]
	F <- 2*p*(1-p)
	sum(F*weights/total.depth)
	}
for(i in 1:nrow(grouse.df3)) grouse.df3$mean.het[i] <- calculateHeterozygosity(grouse.df3[i,])

#exclude low heterozygosity loci? 
writeBayesFst(grouse.df3[grouse.df3$mean.het>0.1,], "../test_files/BB_Fst0_2.inp",gamma=0)

#run bayesfst outside R
source("../test_files/functions.r")

#the following taken from Mark's code
m1 <- data.matrix(read.table("../test_files/BB_Fst0_2.out"))
#read in two functions that will get pvals, fst per locus, and fst per population.

pvals <- getpvals(m1,sum([grouse.df3$mean.het>0.1),3)

#the function getfsts() puts the FST for pops 1..d in columns 1..d. These are samples from the posterior distribution.
fstvals <- getfsts(m1,sum([grouse.df3$mean.het>0.1),15) #nrow(grouse.df3)

plot(pvals[,3],pvals[,1],xlab="transformed p-values",ylab="FST")
#plot critical p-value - in this case the 0.05 level.
x <- 2*abs(0.975 - 0.5)
abline(v=log(x/(1-x)))

#still looking odd

grouse.df4 <- grouse.df3[grouse.df3$mean.het>0.1,]

#try a fisher test on each locus
FisherAlleleTest <- function(data){
	counts <- as.numeric(data[,grep('count',names(data))])
	depth <- as.numeric(data[,grep('depth',names(data))])
	m.table <- cbind(counts,depth-counts)
	fisher.test(m.table)$p
	}

grouse.df3$fisher.pval <- -1
for(i in 1:nrow(grouse.df3)) grouse.df3$fisher.pval[i] <- FisherAlleleTest(grouse.df3[i,])

### fdist2
#het and Fst estimates...
het.fst <- read.table('../test_files/BB_Fst0_2_fst_outfile') #only 1001 loci
names(het.fst) <- c("heterozygosity","Fst")
fdist2.quantiles <- read.table('../test_files/BB_Fst0_2_cplot.out') #95%
names(fdist2.quantiles) <- c("heterozygosity","q02_5","median","q97_5")

plot(Fst~heterozygosity,het.fst)
lines(q02_5~ heterozygosity,fdist2.quantiles)
lines(q97_5~ heterozygosity,fdist2.quantiles)
#looks like crap

#try obtaining distribution of each locus Fst by simulation 
#see functions below




simFst <- function(counts,depth,Fst=0.13,freq=NULL){
	#calculate random beta binomial distributions given Fst
	if(Fst<=0 | Fst>=1) stop('Fst must be in range 0 < Fst < 1\n')
	if(length(counts)!=length(depth)) stop('counts and depth must be equal length\n')
	if(is.null(freq)) freq <- counts/depth
	p.hat <- sum(freq*depth)/sum(depth)
	n.pop <- length(counts)
	
	theta <- 1/Fst -1
	alpha <- theta*p.hat
	beta <- theta*p.hat
	
	p.beta <- rbeta(n=n.pop,alpha,beta)
	
	x.obs <- rbinom(n.pop,size=depth,p=p.beta)
	p.obs <- sum(x.obs)/sum(depth)
	y.obs <- depth-x.obs
	Lik.bin <- numeric(n.pop)
	for(i in 1:n.pop){
		Lik.bin[i] <- lchoose(depth[i],x.obs[i])+(log(p.obs)*x.obs[i])+(log(1-p.obs)*y.obs[i])
		}
	Sigma.Lik.bin <- sum(Lik.bin)
	Sigma.Lik.bin	
	}

FstSimTest <- function(counts,depth,Fst=0.13,freq=NULL,nsim = 1000){
	if(is.null(freq)) freq <- counts/depth
	sim.Lik.bin <- numeric(nsim)
	n.pop <- length(counts)
	p.hat <- sum(freq*depth)/sum(depth)
	
	Lik.bin.obs <- numeric(n.pop)
	for(i in 1:n.pop){
		Lik.bin.obs[i] <- lchoose(depth[i],counts[i])+(log(p.hat)*counts[i])+(log(1-p.hat)*(depth[i]-counts[i]))
		}
	Sigma.Lik.bin.obs <- sum(Lik.bin.obs) #test statistic

	#now generate null distribution given Fst
	
	for(ns in 1:nsim) sim.Lik.bin[ns] <- simFst(counts,depth,Fst)
	tmp.inequality <- sim.Lik.bin<Sigma.Lik.bin.obs
	nonfinite.tests <- !is.finite(tmp.inequality)
	
	out.vec <- c(Sigma.Lik.bin.obs,(sum(tmp.inequality,na.rm=TRUE)+sum(nonfinite.tests))/nsim)
	names(out.vec) <- c("test.stat","sim.pval")
	out.vec
	}

grouse3.FstSim <- data.frame(test.stat =rep(-10000,nrow(grouse.df3)),sim.pval=rep(-10000,nrow(grouse.df3)))

for(snp in 1:nrow(grouse3.FstSim)){
	grouse3.FstSim[snp,] <- FstSimTest(counts=as.numeric(grouse.df3[snp,grep('count',names(grouse.df3))]),depth=as.numeric(grouse.df3[snp,grep('depth',names(grouse.df3))]),Fst=0.13,nsim=1000)
	}

##try to work out Fst more rationally

listFstSim <- vector('list',10)

for(ifst in 1:10){
	listFstSim[[ifst]] <- data.frame(test.stat =rep(-10000,nrow(grouse.df3)),sim.pval=rep(-10000,nrow(grouse.df3)))
	fst.tst <- ifst/40
	for(snp in 1:nrow(grouse.df3)){
		listFstSim[[ifst]][snp,] <- FstSimTest(counts=as.numeric(grouse.df3[snp,grep('count',names(grouse.df3))]),depth=as.numeric(grouse.df3[snp,grep('depth',names(grouse.df3))]),Fst=fst.tst,nsim=1000)
		}
	}

#calculate log likelihood for beta-binomial
LikBetaBinomial <- function(depth,count,p.hat,theta){
	n=as.numeric(depth)
	k=as.numeric(count)
	alpha <- p.hat*theta
	beta <- (1-p.hat)*theta
	logBBlik <- lgamma(n+1) -lgamma(k+1) - lgamma(n-k+1) +
		lgamma(alpha+k) + lgamma(n+beta-k) - lgamma(alpha+beta+n) +
		lgamma(alpha + beta) - lgamma(alpha) - lgamma(beta)
	logBBlik
	}

LikBetaBinomialPops <- function(depth,count,p.hat,theta){
	#run LikBetaBinomial across 3 populations
	n.pops <- length(depth)
	
	logBBlikPops <- 0
	for(i in 1:n.pops){
		logBBlikPops <- logBBlikPops + LikBetaBinomial(depth[i],count[i],p.hat,theta)
		}
	logBBlikPops
	}

#calculate best fit theta for whole dataset

ifst <- seq(from=0.001,to=0.99,length.out=100)
itheta <- 1/ifst -1
lik.theta.fit <- matrix(0,nrow=nrow(grouse.df3),ncol=100,dimnames=list(grouse.df3$id,paste("theta",itheta,sep="_")))


for(i in 1:100){
	tst <- numeric(nrow(grouse.df3))
		for(j in 1:length(tst)){
		lik.theta.fit[j,i] <- LikBetaBinomialPops(depth=as.numeric(grouse.df3[j,grep('depth',names(grouse.df3))]),
		count=as.numeric(grouse.df3[j,grep('count',names(grouse.df3))]),
		p.hat=as.numeric(grouse.df3[j,'p.hat']),theta=itheta[i])
		}
	#lik.theta.fit[i] <- sum(tst)
	}
null.which <- which(colSums(lik.theta.fit)==max(colSums(lik.theta.fit))) #11

BBlikTest.df <- data.frame(null.lik = lik.theta.fit[,null.which],max.lik = rep(0,nrow(grouse.df3)))
BBlikTest.df$max.lik <- apply(lik.theta.fit,1,max)
BBlikTest.df$chi.sq <- 2*(BBlikTest.df$max.lik-BBlikTest.df$null.lik)
BBlikTest.df$p.val <- 1-pchisq(BBlikTest.df$chi.sq,1)

tmp <-  mt.rawp2adjp(BBlikTest.df$p.val,'Bonferroni')
BBlikTest.df$bonferroni <- tmp$adjp[order(tmp$index),'Bonferroni']

tmp <-  mt.rawp2adjp(BBlikTest.df$p.val,'Hochberg')
BBlikTest.df$hochberg <- tmp$adjp[order(tmp$index),'Hochberg']

#repeat with a dodgy method to pool snps at each locus
BBlikTest.pooled <- aggregate(lik.theta.fit,by=list(grouse.df3$cDNA),sum)
BBlikTest.pooled.df <- data.frame(cDNA=BBlikTest.pooled[,1],null.lik = BBlikTest.pooled[,null.which+1],
	max.lik = rep(0,nrow(BBlikTest.pooled)))
BBlikTest.pooled.df$cDNA <- as.character(BBlikTest.pooled.df$cDNA)	
BBlikTest.pooled.df$max.lik <- apply(BBlikTest.pooled[,-1],1,max)
BBlikTest.pooled.df$chi.sq <- 2*(BBlikTest.pooled.df$max.lik-BBlikTest.pooled.df$null.lik)
BBlikTest.pooled.df$p.val <- 1-pchisq(BBlikTest.pooled.df$chi.sq,1)

tmp <-  mt.rawp2adjp(BBlikTest.pooled.df$p.val,'Bonferroni')
BBlikTest.pooled.df$bonferroni <- tmp$adjp[order(tmp$index),'Bonferroni']
tmp <-  mt.rawp2adjp(BBlikTest.pooled.df$p.val,'Hochberg')
BBlikTest.pooled.df$hochberg <- tmp$adjp[order(tmp$index),'Hochberg']
BBlikTest.pooled.df$fst <- 0

BBlikTest.pooled.df$fst <- apply(BBlikTest.pooled[,-1],1,function(X){
	fst.seq <- seq(from=0.001,to=0.99,length.out=100)
	fst.seq[which.max(X)]
	})

BBlikTestSigGenes <- grouse.df3[grouse.df3$cDNA %in% BBlikTest.pooled.df[BBlikTest.pooled.df$hochberg<0.05,]$cDNA,]

grouse.chick.map[grouse.chick.map$grouse.id %in% unique(BBlikTestSigGenes$cDNA),]
chicken.mart <- useMart('ensembl',dataset='ggallus_gene_ensembl')
BBsigBM <- getBM(attributes=c('ensembl_peptide_id','ensembl_gene_id','description'),filters='ensembl_peptide_id',values=grouse.chick.map[grouse.chick.map$grouse.id %in% unique(BBlikTestSigGenes$cDNA),1],mart=chicken.mart)


#add whether immune or not
BBlikTest.pooled.df$type <- factor("control",levels=c("control","immune"))
BBlikTest.pooled.df$type[BBlikTest.pooled.df$cDNA %in% grouse.imm.ids$id] <- "immune"
BBlikTestSigGenes$type <- factor("control",levels=c("control","immune"))
BBlikTestSigGenes$type[BBlikTestSigGenes$cDNA %in% grouse.imm.ids$id] <- "immune"
#need to reannotate this, these aren't picking out all immune genes, eg CD74

BBlikTest.pooled.df2 <- merge(BBlikTest.pooled.df,grouse.chick.map,by.x="cDNA",by.y="grouse.id")
#"cDNA_1760-1" "ssh_1675-1" gone awol, no blast hit?

chick.go_bp <- getBM(attributes=c("ensembl_transcript_id","go_biological_process_id"),filters="ensembl_transcript_id",values=BBlikTest.pooled.df2$ensembl_transcript_id,mart=chicken.mart) 

chick.go_bp <- chick.go_bp[grep('GO',chick.go_bp[,2]),]
library(GO.db)
getAncestorGO <- function(data,envir=GOBPANCESTOR){
	require(GO.db)
	#assume first column gene name and second is GO term
	
	data.out <- data.frame(character(0),character(0),stringsAsFactors=FALSE)
	names(data.out) <- names(data)
	for(gene in unique(data[,1])){
		tmp.df <- data[data[,1]==gene,]
		tmp.GO <- mget(unique(tmp.df[,2]),envir=envir,ifnotfound=NA)
		tmp.GO <- unlist(tmp.GO)
		tmp.GO <- tmp.GO[!is.na(tmp.GO)]
		tmp.GO <- tmp.GO[grep('GO',tmp.GO)]
		tmp.GO <- unique(tmp.GO)
		if(length(tmp.GO)==0) next
		
		tmp.df.out <- data.frame(rep(gene,length(tmp.GO)),tmp.GO,stringsAsFactors=FALSE)
		names(tmp.df.out) <- names(data)
		data.out <- rbind(data.out,tmp.df.out)
		}
	data.out
	}
chick.go_bp2 <- getAncestorGO(chick.go_bp)

chickgenes.with.go <- BBlikTest.pooled.df2[BBlikTest.pooled.df2$ensembl_transcript_id %in% chick.go_bp$ensembl_transcript_id,]
chickgenes.with.go <- chickgenes.with.go[match(unique(chickgenes.with.go$ensembl_transcript_id),chickgenes.with.go$ensembl_transcript_id),]
#names(chickgenes.with.go) <- c("gene_id","p")
chick.geneidx <- create.geneidx(chick.go_bp,go_id="go_biological_process_id",gene_id="ensembl_transcript_id")
chick.geneidx2 <- create.geneidx(chick.go_bp2,go_id="go_biological_process_id",gene_id="ensembl_transcript_id")
chick.geneidx2 <- chick.geneidx2[-1] #first GO term includes all genes

fst.Fang.test <- FangRankTest(geneidx=chick.geneidx,genep=chickgenes.with.go,genep.names=c("ensembl_transcript_id","p.val"))

fst.Fang.test2 <- FangRankTest(geneidx=chick.geneidx2,genep=chickgenes.with.go,genep.names=c("ensembl_transcript_id","p.val"))
#p-values would disappear after correction, biosysnthesis but not immune function


chick.desc <- getBM(attributes=c("ensembl_transcript_id","description"),filters="ensembl_transcript_id",values=grouse.chick.map$ensembl_transcript_id,mart=chicken.mart) 